package main

import (
	"encoding/base64"
	"fmt"
	"io"
	"os"
	"os/exec"
	"strings"
	"syscall"
	"unsafe"

	"github.com/gliderlabs/ssh"
	"github.com/kr/pty"
	gossh "golang.org/x/crypto/ssh"
)

type SshServer struct {
	Port    string
	PemPath string
	auth    Auth
	srv     *ssh.Server
}

func NewSshServer(port, pem, keys string) (srv *SshServer, err error) {
	if _, err = os.Stat(pem); err != nil {
		return
	}
	if _, err = os.Stat(keys); err != nil {
		return
	}
	auth, err := NewLocalAuth(keys)
	if err != nil {
		return
	}
	return &SshServer{
		Port:    port,
		PemPath: pem,
		auth:    auth,
	}, nil
}

func setWinsize(f *os.File, w, h int) { //调整term窗口大小使用 不需要动
	syscall.Syscall(syscall.SYS_IOCTL, f.Fd(), uintptr(syscall.TIOCSWINSZ),
		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
}

func chop(addr string) string {
	if strings.Contains(addr, ":") {
		s := strings.Split(addr, ":")
		return s[0]
	}
	return addr
}

func (ss *SshServer) publicKeyAuth(ctx ssh.Context, key ssh.PublicKey) bool { //使用公钥登陆的验证
	addr := chop(ctx.RemoteAddr().String())
	user := ctx.User()
	k := base64.StdEncoding.EncodeToString(gossh.MarshalAuthorizedKey(key))
	ctx.SetValue("x-public-key", k)
	ok := ss.auth.AuthKey(user, addr, k)
	log.Infof("%s try to login with key, %v", user, ok)
	return ok
}

func (ss *SshServer) passwdAuth(ctx ssh.Context, password string) bool { //使用密码登陆的验证
	user := ctx.User()
	addr := ctx.RemoteAddr().String()
	sarr := strings.Split(addr, ":")
	ok := ss.auth.AuthUserPass(user, password, sarr[0])
	//	ok := ss.auth.AuthUserPass(user, password, "132.121.94.182")
	log.Infof("%s try to login with pass, %v", user, ok)
	return ok
}

func (ss *SshServer) serve() error {
	ssh.Handle(ss.handle)
	ss.srv = &ssh.Server{Addr: ss.Port, Handler: nil}
	ss.srv.SetOption(ssh.HostKeyFile(ss.PemPath))
	ss.srv.SetOption(ssh.PublicKeyAuth(ss.publicKeyAuth)) //使用密钥登陆的设置
	ss.srv.SetOption(ssh.PasswordAuth(ss.passwdAuth))     //使用密码登陆的设置
	log.Infof("start ssh server")
	return ss.srv.ListenAndServe()
}

func (ss *SshServer) handle(session ssh.Session) {
	input := session.Command()
	addr := chop(session.RemoteAddr().String())
	user := session.User()
	log.Infof("handle request %s from %s", user, addr)
	switch {
	case len(input) == 1 && input[0] == "add-key":
		log.Infof("handle %s request %s from %s", "add key", user, addr)
		k := session.Context().Value("x-public-key")
		if k == nil {
			io.WriteString(session, "key is null\n")
			log.Infof("request err is: %s, %s from %s", "key is null", user, addr)
			session.Exit(2)
			break
		}
		key := k.(string)
		if err := ss.auth.AddKey(user, addr, key); err != nil {
			io.WriteString(session, "\n"+err.Error()+"\n")
			log.Infof("request err is: %v, %s from %s", err, user, addr)
			session.Exit(2)
			break
		}
		if err := ss.auth.Unload(); err != nil {
			io.WriteString(session, "\n"+err.Error()+"\n")
			log.Infof("request err is: %v, %s from %s", err, user, addr)
			session.Exit(2)
			break
		}
		io.WriteString(session, "added\n")
		log.Infof("request result is %s, %s from %s", "added", user, addr)
		session.Exit(0)
	case len(input) == 3 && input[0] == "change-pass":
		log.Infof("handle %s request %s from %s", "change pass", user, addr)
		if err := ss.auth.ChangePassTo(user, addr, input[1], input[2]); err != nil {
			io.WriteString(session, fmt.Sprintf("change pass error: %v\n", err))
			log.Infof("request err is: %v, %s from %s", err, user, addr)
			session.Exit(2)
			break
		} else {
			io.WriteString(session, "changed\n")
			log.Infof("request result is %s, %s from %s", "changed", user, addr)
			session.Exit(0)
		}
	default:
		log.Infof("handle %s request %s from %s", "bash", user, addr)
		ptyReq, winCh, isPty := session.Pty()
		if !isPty {
			io.WriteString(session, "No PTY requested\n try to use:\n ssh -tt xx.xx.xx.xx command")
			session.Exit(2)
			break
		}
		cmd := exec.Command("bash")
		cmd.Env = append(cmd.Env, os.Environ()...)
		cmd.Env = append(cmd.Env, fmt.Sprintf("TERM=%s", ptyReq.Term))
		f, err := pty.Start(cmd)
		if err != nil {
			io.WriteString(session, fmt.Sprintf("command error: %v", err))
			log.Infof("request err is: %v, %s from %s", err, user, addr)
			session.Exit(2)
			break
		}
		go func() { //调整窗口大小
			for win := range winCh {
				setWinsize(f, win.Width, win.Height)
			}
		}()
		go func() { //读输入
			io.Copy(f, session) // stdin
		}()
		io.Copy(session, f)
		return
	}
}
